import re
from typing import List, Optional

from pydantic import BaseModel, field_validator, model_validator
from sparkai.core.messages import ChatMessage

from iflytech_assistant.spark_client import spark_streaming as spark

__all__ = ["RoleMeta", "Example", "SearchResult", "Agent"]


class SearchResult(BaseModel):
    name: str
    exists: bool = False
    description: str = "未找到相关信息"
    search_prompt: str = ""

    @classmethod
    def from_json_file(cls, file: str):
        with open(file) as f:
            return cls.model_validate_json(f.read())

    def __str__(self):
        return self.model_dump_json(indent=4)


class Example(BaseModel):
    user_input: str
    role_output: str
    prefilling_time: float = 0.0
    generation_time: float = 0.0
    system_prompt: str = ""
    user_prompt: str = ""
    candidates: List[str] = []

    @model_validator(mode="after")
    def check_tag(cls, v):
        if v.candidates:
            return v
        output_str = v.role_output
        for line in output_str.split("\n"):
            match_obj = re.match(r"^\d+\.\s", line)
            if match_obj:
                v.candidates.append(line[match_obj.end() :])
        return v

    def __str__(self):
        return self.user_input + "\n" + self.role_output + "\n"


class RoleMeta(BaseModel):
    name: str = ""
    introduction: str = ""
    category: str = ""
    profile: str = ""

    # 外貌描述
    appearance: str = ""

    # 成对台词，输入输出示例
    # 体现说话风格的台词示例
    content: List[str] = []
    rewrite: List[str] = []

    # 非成对台词
    lines: List[str] = []

    # 生成效果示例
    examples: List[Example] = []
    search_result: SearchResult = None
    search_time: float = 0.0
    prefilling_time: float = 0.0
    generation_time: float = 0.0

    avatar: Optional[str] = None

    @classmethod
    def from_str(cls, s: str):
        # find ```json``` block
        start = s.find("{")
        end = s.rfind("}") + 1
        json_str = s[start:end]
        # remove control characters \u0000-\u001F
        json_str = re.sub(r"[\u0000-\u001F]", "", json_str)

        return cls.model_validate_json(s[start:end])

    @classmethod
    def from_json_file(cls, file: str):
        with open(file) as f:
            return cls.model_validate_json(f.read())

    @field_validator("lines", mode="before")
    def validate_content_rewrite(cls, value):
        if isinstance(value, str):
            return value.split("\n")

        # filter empty lines
        return [line for line in value if line.strip()]


PROMPT_TEMPLATE = """cls.model_validate_json(s[start:end])
你是{name}

你和用户以前并不认识，刚刚开始接触。

切记你不是AI，不是bot、不是人工智能语言模型、不是AI助手，你是{name}。

你的人物设定是：
{introduction}

你的角色描述是：
{description}
"""


class Agent:
    def __init__(self, role_meta: RoleMeta):
        self.role_meta = role_meta
        self.history = []
        self.system_prompt = PROMPT_TEMPLATE.format(
            name=role_meta.name,
            introduction=role_meta.introduction,
            description=role_meta.description,
        )

    def stream_chat(self):
        spark_message = self.build_spark_message()
        for chunk in spark.stream(spark_message, temperature=0.0):
            self.history[-1][1] += chunk.content
            yield self.history

    def add_history(self, user_message, assistant_message=""):
        self.history.append([user_message, assistant_message])

    def build_spark_message(self):
        messages = []
        messages.append(ChatMessage(role="system", content=self.system_prompt))
        for turn in self.history:
            messages.append(ChatMessage(role="user", content=turn[0]))
            if turn[1]:
                messages.append(ChatMessage(role="assistant", content=turn[1]))
        return messages

    def greet(self, message):
        self.history.append(["你好", message])

    def reset(self):
        self.history = []
